# -*- coding: utf-8 -*-
'''
Created on 24.06.2019

@author: yu03
'''


from FFT_Interpolation import *
from Interpolation_Test import x, tau0, N
from scipy import signal


f = 10.4
phi = 0

sig_origin = np.cos(2*np.pi*f*x + phi) + 1.5
window_gaussian = signal.gaussian(N, std=N/4)
sig = sig_origin * window_gaussian
zero_padd = np.zeros(4096-N)
sig = np.concatenate((sig,zero_padd))
print(len(sig))
# sig = [i+np.random.rand()/10 for i in sig]
# sig = sig / signal.gaussian(N, std=N/4)


figure_name = 'Signal'
plt.figure(figure_name)
plt.gcf().set_size_inches(12,12)
  
plt.subplot(2,1,1)
plt.plot(window_gaussian, label='Gaussian Window', color='red')
plt.plot(sig_origin, label='Original Signal', color='blue')
legend = plt.legend(loc='upper right')
plt.setp(legend.get_texts()[1], color = 'blue')
plt.setp(legend.get_texts()[0], color = 'red')
plt.title("Signal Cutting")
plt.ylabel("Amplitude")
plt.xlabel("Samples")
plt.grid(which='major', axis='both')
  
plt.subplot(2,1,2)
plt.plot(sig, label='Signal', color='blue')
plt.title("Signal")
plt.ylabel("Amplitude")
plt.xlabel("Frequency")
plt.grid(which='major', axis='both')
  
plt.tight_layout()
plt.show()

# plt.savefig('%s.png'%figure_name, dpi=300)

# sig = sig - np.average(sig)
freqline, sig_FFT, sig_freq, sig_phase = FFT_cal(sig, tau0)

freq_bin = 1/ tau0 / N
m_k_num = sig_freq[10:].argmax() + 10 ### frequency bin 序数
m_k = freqline[m_k_num] ### 换算为实际频率
X_m_k = sig_FFT[m_k_num] ### magnitude (complex raw data)
X_m_k_minus = sig_FFT[m_k_num-1]
X_m_k_plus = sig_FFT[m_k_num+1]
C_1 = (X_m_k_minus - X_m_k_plus) / (2 * X_m_k - X_m_k_minus - X_m_k_plus)
freq_estim_1 = np.real(C_1) * freq_bin + m_k ### No windowing
print(freq_estim_1)

N_sel = int(N * int(freq_estim_1) / freq_estim_1)
print(N_sel)
sig_sel = sig[:N_sel]
sig_sel = sig_sel - np.average(sig_sel)
freqline_sel, sig_FFT_sel, sig_freq_sel, sig_phase_sel = FFT_cal(sig_sel, tau0)

figure_name = 'Frequency Estimation'
plt.figure(figure_name)
plt.gcf().set_size_inches(12,12)

plt.subplot(3,1,1)
plt.plot(sig, label='Original Signal', color='blue')
plt.plot(sig_sel, label='Selected Signal', color='red')
legend = plt.legend(loc='lower right')
plt.setp(legend.get_texts()[0], color = 'blue')
plt.setp(legend.get_texts()[1], color = 'red')
plt.title("Signal Cutting")
plt.ylabel("Amplitude")
plt.xlabel("Samples")
plt.grid(which='major', axis='both')
plt.subplot(3,1,2)
plt.stem(freqline, sig_freq, linefmt='b', markerfmt='bo', basefmt=" ", label='Magnitude Spectrum')
plt.plot([freq_estim_1,freq_estim_1], [0,1], label='Estimated Freq.', color='red')
legend = plt.legend(loc='upper right')
plt.setp(legend.get_texts()[1], color = 'blue')
plt.setp(legend.get_texts()[0], color = 'red')
plt.title("Freq. Estimation")
plt.ylabel("Magnitude")
plt.xlabel("Frequency")
plt.xlim(0,30)
plt.grid(which='major', axis='both')


plt.subplot(3,1,3)
plt.stem(freqline_sel, sig_freq_sel, linefmt='b', markerfmt='bo', basefmt=" ")
plt.title("Recalculated FFT")
plt.ylabel("Magnitude")
plt.xlabel("Frequency")
plt.xlim(0,30)
plt.grid(which='major', axis='both')
plt.tight_layout()
plt.show()


